import numpy as np
from collections import deque

class TangentForceRateCalculator:
    """Calculate the rate of change of tangent force."""
    def __init__(self, window_size=5, momentum=0.9):
        """
        Initialize the calculator with a window size.
        Args:
            window_size (int): The size of the window for calculating the rate.
        """
        self.window_size = window_size
        self.previous_tangent_force = deque(maxlen=window_size)
        self.momentum = momentum
        self.now_rate = 0.0
        x = np.arange(window_size)
        self.A = np.vstack([x, np.ones(len(x))]).T
    
    def reset(self):
        """
        Reset the calculator.
        """
        self.previous_tangent_force.clear()
        self.now_rate = 0.0
    
    def calculate_rate(self, tangent_force : float):
        """
        Calculate the rate of change of tangent force.
        Args:
            tangent_force (float): The current tangent force value.
        Returns:
            float: The rate of change of tangent force.
        """

        self.previous_tangent_force.append(tangent_force)
        if len(self.previous_tangent_force) < self.window_size:
            return 0.0
        
        # Calculate the rate of change with all the data in the window using linear fitting

        y = np.array(self.previous_tangent_force)

        m, c = np.linalg.lstsq(self.A, y, rcond=None)[0]
        self.now_rate = self.now_rate * self.momentum + m * (1 - self.momentum)
        return self.now_rate

class StateManager:
    def __init__(self, SN1, SN2):
        self.contact_area = {SN1: None, SN2: None}
        self.contact_center = {SN1: None, SN2: None}
        self.tangential_force = {SN1: None, SN2: None}
        self.tangential_force_rate = {SN1: None, SN2: None}
        self.rate_calculator = {SN1: TangentForceRateCalculator(window_size=5, momentum=0.9),
                                SN2: TangentForceRateCalculator(window_size=5, momentum=0.9)}
    
    def reset(self):
        """
        Resets the state manager, clearing all contact areas, centers, and tangential forces.
        """
        self.contact_area = {k: None for k in self.contact_area.keys()}
        self.contact_center = {k: None for k in self.contact_center.keys()}
        self.tangential_force = {k: None for k in self.tangential_force.keys()}
        self.tangential_force_rate = {k: None for k in self.tangential_force_rate.keys()}
        self.rate_calculator = {k: v.reset() for k,v in self.rate_calculator.items()}

    def update_contact_area(self, SN, F):
        """
        Updates the contact area for a given sensor SN based on the force data F.

        Args:
            SN (str): Sensor ID.
            F (np.ndarray): Force data.
        """
        mask = np.linalg.norm(F, axis=1) > 0.02
        self.contact_area[SN] = np.sum(mask)

    def update_contact_center(self, SN, F):
        contact_mask = np.linalg.norm(F, axis=1) > 0.01
        contact_mask = contact_mask.reshape(20,20)
        contact_positions = np.stack(np.where(contact_mask))
        if contact_positions.size == 0:
            self.contact_center[SN] = None
            return
        self.contact_center[SN] = np.mean(contact_positions, axis=1)
    
    def update_tangential_force(self, SN, F):
        """
        Updates the tangential force for a given sensor SN based on the force data F.

        Args:
            SN (str): Sensor ID.
            F (np.ndarray): Force data.
        """
        self.tangential_force[SN] = np.sum(np.linalg.norm(F[:, :2], axis=1))
        self.tangential_force_rate[SN] = self.rate_calculator[SN].calculate_rate(self.tangential_force[SN])
    
    @property
    def avg_contact_area(self):
        """
        Computes the average contact area across all sensors.

        Returns:
            float: Average contact area.
        """
        return np.mean([v for v in self.contact_area.values() if v is not None])
    
    @property
    def avg_tangential_force(self):
        """
        Computes the average tangential force across all sensors.

        Returns:
            float: Average tangential force.
        """
        return np.mean([v for v in self.tangential_force.values() if v is not None])
    
    @property
    def avg_tangential_force_rate(self):
        """
        Computes the average tangential force rate across all sensors.

        Returns:
            float: Average tangential force rate.
        """
        return np.mean([v for v in self.tangential_force_rate.values() if v is not None])